#!/usr/bin/env python
# coding: utf-8

# ## 导入库

# In[1]:

from collections import defaultdict

from machine_lib import *

# ## 1, 登录
# 获取会话

# 1, 去machine_lib.py的login()函数中填写用户名和密码
# 2, 保存修改后的machine_lib.py文件
# 3, 返回本文件，重启内核
# 4, 重新运行第一格代码：from machine_lib import *
# 5, 所有machine_lib中的改动都要执行上述1-4步骤才能生效

# In[2]:

session = load_session()

# ## 2, 获取数据字段

# In[3]:

df = get_datafields(session, dataset_id='analyst4', region='USA', universe='TOP3000', delay=1)
print(df)

# In[4]:

print(df[df['type'] == "MATRIX"]["id"].tolist())
pc_fields = process_datafields(df, "matrix")
print(pc_fields)

# ## 3, Alpha工厂

# In[5]:
with open('params.json', 'r') as f:
    params = json.load(f)
fields = params.get('fields', [])

# 调用 get_first_order 函数，并传递 fields 作为第三个参数
first_order = get_first_order(pc_fields, ts_ops + basic_ops, fields)
print(first_order[:10])
print(len(first_order))

# In[6]:

# 加载区域和衰减系数的alpha

fo_alpha_dict = defaultdict(list)
init_decay = 6
regions = ["usa"]
for key in regions:
    for alpha in first_order:
        fo_alpha_dict[key].append((alpha, init_decay))

for key, value in fo_alpha_dict.items():
    print("%s : %d" % (key, len(value)))

# ## 4, 模拟alpha

# In[7]:

region_dict = {"usa": ("USA", "TOP3000"), "asi": ("ASI", "MINVOL1M"), "eur": ("EUR", "TOP1200"),
               "glb": ("GLB", "TOP3000"), "hkg": ("HKG", "TOP800"), "twn": ("TWN", "TOP500"), "jpn": ("JPN", "TOP1600"),
               "kor": ("KOR", "TOP600"), "chn": ("CHN", "TOP2000U"), "amr": ("AMR", "TOP600")}

# In[10]:
first_layer_bag = []
max_count = 20000  # 设置最大计数值

while get_current_count() < max_count:
    try:
        simulate(fo_alpha_dict, region_dict, "analyst4_fo_usa", 'SUBINDUSTRY', get_current_count(), first_layer_bag)
        break  # 如果simulate成功执行，没有异常抛出，跳出循环
    except Exception as e:
        print(f"An error occurred: {e}. Retrying...")
        login()
        if get_current_count() >= max_count:
            print(f"CURRENT_COUNT reached {max_count}. Exiting...")
            break

# ## 5, 选择alpha
# 转到网络alpha面板，查找编号和日期以跟踪下一个顺序的改进

# In[ ]:

## 获取有希望的alpha以在下一个顺序中进行改进
fo_tracker = get_alphas("08-25", "08-26", 1.2, 0.5, "USA", 500, "track")

# In[ ]:

print(len(fo_tracker['next']))
print(len(fo_tracker['decay']))
fo_layer = transform(fo_tracker['next'] + fo_tracker['decay'], 'usa')

# ## 5, 下一步改进
# 二阶：ts_ops(field, days) -> group_ops(ts_ops(field, days), group)

# In[ ]:

so_alpha_dict = defaultdict(list)
group_ops = ["group_neutralize", "group_rank", "group_normalize", "group_scale", "group_zscore"]
for key, couples in fo_layer.items():
    for expr, decay in couples:
        for alpha in get_group_second_order_factory([expr], group_ops, key):
            so_alpha_dict[key].append((alpha, decay))

for key, value in so_alpha_dict.items():
    print("%s : %d" % (key, len(value)))

print(so_alpha_dict['usa'][:3])

# In[ ]:

# 模拟二阶alpha
so_bag = []
simulate(so_alpha_dict, region_dict, "oth432_usa_fo", 'SUBINDUSTRY', 0, so_bag)

# ## 7, 获取可提交的alpha
# 检查提交并查看alpha

# In[ ]:

# 获取可提交的alpha并检查提交
th_tracker = get_alphas("07-16", "07-20", 1.58, 1, "oth432_usa_so", 200, "submit")

# In[ ]:

stone_bag = []
for alpha in th_tracker['next'] + th_tracker['decay']:
    stone_bag.append(alpha[0])
print(len(stone_bag))
gold_bag = []
check_submission(stone_bag, gold_bag, 0, environ['consultant_username'])

# In[ ]:

# 查看日期和指标以定位网络中的alpha
view_alphas(gold_bag)

# ## 8, 微调可提交的alpha
# 中性化、性能比较、周转、利润率

# In[ ]:
